import torch
import math
import numpy as np
import matplotlib.pyplot as plt


def log_sinkhorn_iterations(Z, log_mu, log_nu, iters: int):
    """ Perform Sinkhorn Normalization in Log-space for stability"""
    u, v = torch.zeros_like(log_mu), torch.zeros_like(log_nu)
    for _ in range(iters):
        u = log_mu - torch.logsumexp(Z + v.unsqueeze(1), dim=2)
        v = log_nu - torch.logsumexp(Z + u.unsqueeze(2), dim=1)
    return Z + u.unsqueeze(2) + v.unsqueeze(1)

def log_optimal_transport(scores, alpha, iters: int):
    """ Perform Differentiable Optimal Transport in Log-space for stability"""
    b, m, n = scores.shape
    one = scores.new_tensor(1)
    ms, ns = (m*one).to(scores), (n*one).to(scores)

    bins0 = alpha.expand(b, m, 1)
    bins1 = alpha.expand(b, 1, n)
    alpha = alpha.expand(b, 1, 1)

    couplings = torch.cat([torch.cat([scores, bins0], -1),
                           torch.cat([bins1, alpha], -1)], 1)

    norm = - (ms + ns).log()
    log_mu = torch.cat([norm.expand(m), ns.log()[None] + norm])
    log_nu = torch.cat([norm.expand(n), ms.log()[None] + norm])
    log_mu, log_nu = log_mu[None].expand(b, -1), log_nu[None].expand(b, -1)

    Z = log_sinkhorn_iterations(couplings, log_mu, log_nu, iters)
    Z = Z - norm  # multiply probabilities by M+N
    return Z

def get_match_indice(scores, match_threshold=0.0):

    def arange_like(x, dim: int):
        return x.new_ones(x.shape[dim]).cumsum(0) - 1  # traceable in 1.1

    # Get the matches with score above "match_threshold".
    max0, max1 = scores[:, :-1, :-1].max(2), scores[:, :-1, :-1].max(1)
    indices0, indices1 = max0.indices, max1.indices
    mutual0 = arange_like(indices0, 1)[None] == indices1.gather(1, indices0)
    mutual1 = arange_like(indices1, 1)[None] == indices0.gather(1, indices1)
    zero = scores.new_tensor(0)
    mscores0 = torch.where(mutual0, max0.values.exp(), zero)
    mscores1 = torch.where(mutual1, mscores0.gather(1, indices1), zero)
    valid0 = mutual0 & (mscores0 > match_threshold)
    valid1 = mutual1 & valid0.gather(1, indices1)
    indices0 = torch.where(valid0, indices0, indices0.new_tensor(-1))
    indices1 = torch.where(valid1, indices1, indices1.new_tensor(-1))

    return {
        'matches0': indices0,  # use -1 for invalid match
        'matches1': indices1,  # use -1 for invalid match
        'matching_scores0': mscores0,
        'matching_scores1': mscores1,
    }


def pdf_normal(n, std, mean):
    f = lambda x: 1/(2*math.pi)**0.5/std\
                                *math.exp(-((x-mean)/std)**2/2)
    data = []
    for i in range(n):
        data.append(f(i))
    return torch.tensor(data)

if __name__ == '__main__':
    img_size = [320, 480]

    n = 100
    x = torch.arange(n) + 1

    px = pdf_normal(n, mean=20, std=5)
    qx = pdf_normal(n, mean=60, std=15)

    # qx = qx * (px.max()/qx.max())
    X, Y = np.meshgrid(px[0:100], qx[0:100])
    # C = -np.abs(X - Y)
    C = -(np.abs(X-Y))
    S = torch.tensor(C, dtype=torch.float32).unsqueeze(0)

    bin_score = torch.tensor(1.)
    scores = log_optimal_transport(S, bin_score, iters=10)
    re = get_match_indice(scores)

    plt.figure()
    # plt.imshow(C, interpolation='none', origin='lower')
    tmp = scores[0,:-1,:-1].exp().numpy()
    plt.imshow(tmp, interpolation='none', origin='lower')
    plt.show()

    print('Done!')